import rawpy
import numpy as np
from PIL import Image
from pathlib import Path

# 尝试导入 pillow-heif，决定 HEIC 是否可用
try:
    from pillow_heif import register_heif_opener
    register_heif_opener()
    HEIC_SUPPORTED = True
except ImportError:
    HEIC_SUPPORTED = False

def extract_raw_thumbnail(raw_path):
    """
    尝试从RAW文件中提取缩略图作为备用方案

    Args:
        raw_path: RAW文件路径

    Returns:
        PIL.Image: 提取的缩略图
    """
    try:
        with rawpy.imread(str(raw_path)) as raw:
            # 尝试提取缩略图
            thumb = raw.extract_thumb()
            if thumb.format == rawpy.ThumbFormat.JPEG:
                # 缩略图是JPEG格式
                from io import BytesIO
                thumb_image = Image.open(BytesIO(thumb.data))
                thumb_image.info['dpi'] = (72, 72)
                return thumb_image
            elif thumb.format == rawpy.ThumbFormat.BITMAP:
                # 缩略图是位图格式
                thumb_image = Image.fromarray(thumb.data)
                thumb_image.info['dpi'] = (72, 72)
                return thumb_image
    except Exception as e:
        print(f"无法提取RAW缩略图: {e}")
        raise


def load_raw_image(raw_path, preview_mode=False):
    """
    加载RAW图像并转换为PIL Image对象

    Args:
        raw_path: RAW文件路径
        preview_mode: 是否使用预览模式（更快但质量较低）

    Returns:
        PIL.Image: 转换后的图像对象
    """
    # 如果是预览模式，优先尝试提取缩略图
    if preview_mode:
        try:
            return extract_raw_thumbnail(raw_path)
        except:
            # 如果缩略图提取失败，继续使用常规方法但使用快速参数
            pass

    # 多次尝试加载，解决"Out of order call"问题
    max_attempts = 3

    for attempt in range(max_attempts):
        try:
            # 每次尝试都重新打开文件，避免状态问题
            with rawpy.imread(str(raw_path)) as raw:
                # 根据模式选择不同的处理参数
                if preview_mode:
                    # 预览模式：使用快速参数
                    processing_attempts = [
                        {
                            'use_camera_wb': True,
                            'half_size': True,  # 使用半尺寸加快速度
                            'no_auto_bright': True,
                            'output_bps': 8
                        },
                        {
                            'half_size': True,
                            'no_auto_bright': True,
                            'output_bps': 8
                        }
                    ]
                else:
                    # 完整模式：使用高质量参数
                    processing_attempts = [
                        # 第一次尝试：基本参数
                        {
                            'use_camera_wb': True,
                            'half_size': False,
                            'no_auto_bright': True,
                            'output_bps': 8
                        },
                        # 第二次尝试：使用自动白平衡
                        {
                            'use_camera_wb': False,
                            'half_size': False,
                            'no_auto_bright': False,
                            'output_bps': 8,
                            'user_wb': [1.0, 1.0, 1.0, 1.0]
                        },
                        # 第三次尝试：使用半尺寸处理（可能有助于某些损坏的文件）
                        {
                            'use_camera_wb': False,
                            'half_size': True,
                            'no_auto_bright': False,
                            'output_bps': 8
                        },
                        # 第四次尝试：最简单的参数
                        {}
                    ]

                rgb = None
                last_error = None

                for i, params in enumerate(processing_attempts):
                    try:
                        rgb = raw.postprocess(**params)
                        break  # 成功就退出循环
                    except Exception as e:
                        last_error = e
                        if i < len(processing_attempts) - 1:
                            continue  # 尝试下一种参数
                        else:
                            # 所有尝试都失败了
                            raise e

                # 转换为PIL Image
                pil_image = Image.fromarray(rgb)

                # 如果是预览模式且使用了半尺寸，需要调整尺寸信息
                if preview_mode and 'half_size' in processing_attempts[0] and processing_attempts[0]['half_size']:
                    # 将尺寸信息翻倍以反映真实尺寸
                    original_size = pil_image.size
                    pil_image.info['original_size'] = (original_size[0] * 2, original_size[1] * 2)

                # 尝试获取EXIF信息
                try:
                    # 从RAW文件中提取基本信息
                    dpi = 72  # RAW文件通常没有DPI信息，使用默认值

                    # 将DPI信息添加到PIL图像
                    pil_image.info['dpi'] = (dpi, dpi)

                except Exception as e:
                    print(f"无法从RAW文件获取EXIF信息: {e}")
                    pil_image.info['dpi'] = (72, 72)

                return pil_image

        except rawpy.LibRawError as e:
            error_str = str(e)

            # 特殊处理"Out of order call"错误
            if "out of order" in error_str.lower():
                if attempt < max_attempts - 1:
                    print(f"RAW文件状态错误，正在重试... (尝试 {attempt + 1}/{max_attempts})")
                    continue  # 重试
                else:
                    # 最后尝试提取缩略图
                    print("尝试提取RAW文件缩略图作为备用方案...")
                    try:
                        return extract_raw_thumbnail(raw_path)
                    except:
                        pass

                    error_msg = f"RAW文件处理失败 (.rw2): 多次尝试后仍然出现状态错误"
                    error_msg += f"\n\n可能的原因:\n"
                    error_msg += f"• 该RAW文件可能存在内部结构问题\n"
                    error_msg += f"• rawpy库版本与该文件格式不完全兼容\n"
                    error_msg += f"• 文件可能在传输过程中损坏\n"
                    error_msg += f"\n建议:\n"
                    error_msg += f"• 尝试使用Panasonic官方软件打开该文件\n"
                    error_msg += f"• 考虑将文件转换为DNG格式\n"
                    error_msg += f"• 检查文件是否完整下载/复制\n"
                    error_msg += f"• 跳过此文件，处理其他文件"
                    raise Exception(error_msg)

            # 其他rawpy错误
            file_ext = Path(raw_path).suffix.lower()
            error_msg = f"RAW文件处理失败 ({file_ext}): {error_str}"
            if "Data error" in error_str or "unsupported" in error_str.lower():
                error_msg += f"\n\n可能的原因:\n"
                error_msg += f"• 该{file_ext}格式不被当前rawpy版本支持\n"
                error_msg += f"• 文件可能损坏或不完整\n"
                error_msg += f"• 相机型号过新，需要更新rawpy库\n"
                error_msg += f"• 该文件可能不是标准的RAW格式"
            raise Exception(error_msg)

        except FileNotFoundError:
            raise Exception(f"RAW文件不存在: {raw_path}")

        except PermissionError:
            raise Exception(f"没有权限访问RAW文件: {raw_path}")

        except MemoryError:
            raise Exception(f"内存不足，无法处理大型RAW文件: {raw_path}")

        except Exception as e:
            # 通用错误处理
            file_ext = Path(raw_path).suffix.lower()
            error_msg = f"加载RAW文件失败 ({file_ext}): {str(e)}"

            # 检查是否是常见的不支持格式错误
            if any(keyword in str(e).lower() for keyword in ['unsupported', 'format', 'decode', 'invalid']):
                error_msg += f"\n\n建议:\n"
                error_msg += f"• 检查文件是否为有效的RAW格式\n"
                error_msg += f"• 尝试使用相机厂商的RAW转换工具\n"
                error_msg += f"• 考虑将RAW文件转换为DNG格式后再处理"

            raise Exception(error_msg)

    # 如果所有尝试都失败了（不应该到达这里）
    raise Exception(f"RAW文件处理失败: 多次尝试后仍无法加载文件")


def safe_open_image(image_path):
    """
    安全地打开各种格式的图像文件，包括HEIC格式
    """
    try:
        return Image.open(image_path)
    except Exception as e:
        # 如果是HEIC格式且pillow-heif未安装，提供错误提示
        if str(image_path).lower().endswith(('.heic', '.heif')) and not HEIC_SUPPORTED:
            raise Exception(f"HEIC格式不受支持。请安装pillow-heif库：pip install pillow-heif")
        else:
            raise e